import torch
import torch.nn as nn
import torch.nn.functional as F
from mmcv.cnn import ConvModule, xavier_init
import math
from mmdet.core import auto_fp16
from torch.nn.modules.utils import _pair
from torch.nn import init
from torchvision.ops.deform_conv import deform_conv2d as deform_conv2d_tv
from torch import Tensor
from mmdet.utils import get_root_logger
from einops import rearrange
from inspect import isfunction
from timm.models.layers import PatchEmbed, Mlp, DropPath, trunc_normal_, lecun_normal_
import numpy as np
from functools import partial


from ..builder import NECKS
BatchNorm2d = nn.BatchNorm2d


@NECKS.register_module()
class TSFI_MOE(nn.Module):

    def __init__(self,
                 in_channels,#
                 out_channels,#
                 num_outs,#
                 conv_cfg=None,
                 norm_cfg=None,
                 act_cfg=None,
                 final_embed_dim = 256,
                 upsample_cfg=dict(mode='nearest')):
        
        super(TSFI_MOE, self).__init__()
        assert isinstance(in_channels, list)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.num_ins = len(in_channels)
        self.num_outs = num_outs
        self.fp16_enabled = False
        self.upsample_cfg = upsample_cfg.copy()

        self.rank_list = [16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128]
        self.pixel_no = [256,128,64,32]
        self.tasks = ['rd', 'ss']
        self.num_layers = len(self.in_channels)
        self.fea_fuse = nn.ModuleList()

        self.final_embed_dim = final_embed_dim
        self.num_lora = 15
        self.MLoRE_1 = nn.ModuleList()
        self.MLoRE_2 = nn.ModuleList()

        for i_layer in range(self.num_layers):
            self.fea_fuse.append(nn.Conv2d(256, final_embed_dim, kernel_size=1, stride=1))
            self.MLoRE_1.append(MOEBlock(final_embed_dim, im_size=self.pixel_no[i_layer]**2, kernel_size=3,with_feat=False))
            self.MLoRE_2.append(MOEBlock(final_embed_dim, im_size=self.pixel_no[i_layer]**2, kernel_size=3,with_feat=False))

        # for il in range(self.num_layers):
        #     self.MLoRE_2.append(MOEBlock(final_embed_dim, im_size=self.pixel_no[i_layer]**2, kernel_size=3,with_feat=False))
        # self.task_mask = nn.ModuleDict()
        
        # for task in self.tasks:
        #     self.task_mask[task] = nn.Sequential(
        #                                         nn.Conv2d(final_embed_dim * 4, final_embed_dim, kernel_size=1),
        #                                         BatchNorm2d(final_embed_dim), nn.GELU(),  
        #                                         nn.Conv2d(final_embed_dim, self.num_layers, kernel_size=3, padding=1))
        # norm_layer = partial(nn.LayerNorm, eps=1e-6)
        # self.norm = norm_layer(256)


    # default init_weights for conv(msra) and norm in ConvModule
    def init_weights(self):
        """Initialize the weights of FPN module"""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                xavier_init(m, distribution='uniform')

    @auto_fp16()
    def forward(self, inputs):
        """Forward function"""
        # logger = get_root_logger()
        # for i in range(len(inputs)):
        #     print(f'inputs[{i}]: {inputs[i].shape}')
        # inputs[0]: torch.Size([2, 256, 256, 256])
        # inputs[1]: torch.Size([2, 256, 128, 128])
        # inputs[2]: torch.Size([2, 256, 64, 64])
        # inputs[3]: torch.Size([2, 256, 32, 32])

        assert len(inputs) == len(self.in_channels)

        all_tasks = self.tasks
        out_feat = {task: [] for task in all_tasks}
        last_feat = {task:0 for task in all_tasks}
        out_mask = {task: 0 for task in all_tasks}
        info = {} # pass information through the pipeline

        x_q_list = {task:[] for task in all_tasks}
        route_feat_1 = {task: None for task in all_tasks}
        route_feat_2 = {task: None for task in all_tasks}
        route_prob_1 = [{task:None for task in all_tasks} for i in range(len(inputs))]
        route_prob_2 = [{task:None for task in all_tasks} for i in range(len(inputs))]

        # if self.num_outs > len(outs):
        #     outs.append(F.max_pool2d(outs[-1], 1, stride=2)) # 对最后一个尺度的特征做maxpooling
        # return ss_feat, rd_feat, cls_feat
        # return tuple(origin_out), ss_feat, rd_feat, cls_feat

        for idx in range(len(inputs)):
            x = inputs[idx]
            # torch.Size([4, 256, 256, 256])
            _cur_task_fea, x_q, route_feat_1, route_prob_1[idx] = self.cal_task_feature(x, idx, route_feat_1)
            # print('done') # _cur_task_fea['rd'].shape torch.Size([2, 320, 256, 256])
            # x_q = _cur_task_fea route_feat_1

            for task in all_tasks:
                _cur_task_fea_now, route_feat_2[task], route_prob_2[idx][task] = self.MLoRE_2[idx](_cur_task_fea[task], task, route_feat_2[task])
                # x_q_list[task].append(x_q[task])
                out_feat[task].append(_cur_task_fea_now)
        
        
        if self.num_outs > len(inputs):
            out_feat['rd'].append(F.max_pool2d(out_feat['rd'][-1], 1, stride=2)) # 对最后
        rd_feat = tuple(out_feat['rd'])
        ss_feat = tuple(out_feat['ss'])
        cls_feat = rd_feat[-1]

        return ss_feat, rd_feat, cls_feat

        # return last_feat, info



    def cal_task_feature(self, x, idx, route_feat):
        ''' Calculate task feature at this layer
        '''
        # combined_fea = rearrange(x, 'b (h w) c -> b c h w', h=self.resolution[0], w=self.resolution[1])

        combined_fea = self.fea_fuse[idx](x) # torch.Size([2, 320, 256, 256])

        x_q = {task:0 for task in self.tasks} # {'rd': 0, 'ss': 0, 'cls': 0}
        route_feat_out = {task:0 for task in self.tasks} # {'rd': 0, 'ss': 0, 'cls': 0}
        route_prob_out = {task:0 for task in self.tasks} # {'rd': 0, 'ss': 0, 'cls': 0}
        for task in self.tasks:
            x_q[task], route_feat_out[task], route_prob_out[task] = self.MLoRE_1[idx](combined_fea, task, route_feat[task])

        combined_fea = x_q

        return combined_fea, x_q, route_feat_out, route_prob_out
       




class Attention(nn.Module):
    def __init__(self, chan_nheads, resolution, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

        self.dim = dim
        self.resolution = resolution
        pixel_no = int(resolution[0] * resolution[1])
        self.pixel_no = pixel_no

        self.chan_nheads = chan_nheads
        chan_head_dim = self.pixel_no // self.chan_nheads
        self.chan_scale = chan_head_dim ** -0.5

    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]   

        raw_spa_attn = (q @ k.transpose(-2, -1))
        attn = raw_spa_attn * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
        raw_spa_attn = raw_spa_attn, attn

        x = (attn @ v).transpose(1, 2).reshape(B, N, C) # (B, task_no+1+HxW, C)

        x = self.proj(x)
        x = self.proj_drop(x)
        raw_attn = [raw_spa_attn]

        return x, raw_attn

class LoraBlock(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size=1, rank=6):
        super().__init__()
        self.W = nn.Conv2d(in_channels, rank, kernel_size=kernel_size, stride=1, padding=kernel_size//2)
        self.M = nn.Conv2d(rank, out_channels, kernel_size=1, stride=1)

    def init_weights(self):
        nn.init.kaiming_uniform_(self.W.weight, a=math.sqrt(5))
        nn.init.zeros_(self.W.bias)
        nn.init.kaiming_uniform_(self.M.weight, a=math.sqrt(5))
        nn.init.zeros_(self.M.bias)
    
    def forward(self, x):
        x = self.W(x)
        x = self.M(x)
        return x

class Block(nn.Module):

    def __init__(self, chan_nheads, resolution, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(chan_nheads, resolution, dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

    def forward(self, x):

        x_attn, attn_weight = self.attn(self.norm1(x))
        x = x + self.drop_path(x_attn)
        x = x + self.drop_path(self.mlp(self.norm2(x)))

        return x, attn_weight

class SpatialAtt(nn.Module):
    def __init__(self, dim, dim_out, im_size, with_feat):
        super().__init__()
        self.conv1 = nn.Conv2d(dim, dim_out, kernel_size=1)
        self.act = nn.GELU()
        self.ln = nn.LayerNorm(dim_out)
        self.convsp = nn.Linear(im_size, 1)
        self.ln_sp = nn.LayerNorm(dim)
        self.conv2 = nn.Conv2d(dim, dim_out, kernel_size=1)
        self.conv3 = nn.Conv2d(dim_out, dim_out, kernel_size=1)
        self.with_feat = with_feat
        if with_feat:
            self.feat_linear = nn.Conv2d(dim_out *2 , dim_out *2, kernel_size=1)
    
    def forward(self, x, route_feat=None):
        n, _, h, w = x.shape # ([2, 320, 256, 256])
        feat = self.conv1(x) # torch.Size([2, 80, 256, 256])
        feat = self.ln(feat.reshape(n, -1, h * w).permute(0, 2, 1)).permute(0, 2, 1).reshape(n, -1, h, w)
        feat = self.act(feat)
        feat = self.conv3(feat)

        feat_sp = self.convsp(x.reshape(n, -1, h * w)).reshape(n, 1, -1)
        feat_sp = self.ln_sp(feat_sp).reshape(n, -1, 1, 1)
        feat_sp = self.act(feat_sp)
        feat_sp = self.conv2(feat_sp)
        
        n, c, h, w = feat.shape
        feat = torch.mean(feat.reshape(n, c, h * w), dim=2).reshape(n, c, 1, 1)
        feat = torch.cat([feat, feat_sp], dim=1)

        return feat

class MOEBlock(nn.Module):
    def __init__(self, final_embed_dim, im_size, kernel_size=1, with_feat=False):
        super().__init__()
        self.rank_list = [16, 24, 32, 40, 48, 56, 64, 72, 80, 88]
        self.spe_rank=64
        self.topk = 9


        self.num_lora = len(self.rank_list)
        self.lora_list_1 = nn.ModuleList()
        for i in range(self.num_lora):
            self.lora_list_1.append(LoraBlock(final_embed_dim, final_embed_dim, kernel_size=kernel_size, rank=self.rank_list[i]))
            self.lora_list_1[i].init_weights()
        self.conv1 = nn.ModuleDict()
        self.conv2 = nn.ModuleDict()
        self.conv3 = nn.ModuleDict()
        self.share_conv = nn.Conv2d(final_embed_dim, final_embed_dim, kernel_size=3, padding=1)
        self.bn = nn.ModuleDict()
        self.bn_all = nn.ModuleDict()
        self.activate = nn.GELU()
        for task in ['rd', 'ss']:
            self.conv1[task] = nn.Conv2d(final_embed_dim, final_embed_dim, kernel_size=1)
            self.conv3[task] = nn.Conv2d(final_embed_dim, final_embed_dim, kernel_size=1)
            self.conv2[task] = LoraBlock(final_embed_dim, final_embed_dim, kernel_size=kernel_size, rank=self.spe_rank)

            self.bn[task] = BatchNorm2d(final_embed_dim)
            self.bn_all[task] = BatchNorm2d(final_embed_dim)

        self.router_1 = nn.ModuleDict() 
        self.desert_k = len(self.rank_list) - self.topk
        for task in ['rd', 'ss']:
            self.router_1[task] = nn.ModuleList()
            self.router_1[task].append(SpatialAtt(final_embed_dim, final_embed_dim // 4, im_size=im_size, with_feat=with_feat))
            self.router_1[task].append(nn.Conv2d(final_embed_dim // 2, self.num_lora * 2 + 1, kernel_size=1))
        
    def forward(self, x, task, route_feat_in=None):
        out_ori = self.conv1[task](x) # torch.Size([2, 320, 256, 256])
        out = out_ori
        n, c, h, w = out.shape
        route_feat = self.router_1[task][0](out, route_feat_in) 
        prob_all = self.router_1[task][1](route_feat).unsqueeze(2) # torch.Size([2, 31, 1, 1, 1])
        prob_lora, prob_mix = prob_all[:, :self.num_lora * 2], prob_all[:, self.num_lora * 2:]
        route_1_raw, stdev_1 = prob_lora.chunk(2, dim=1)  # n, 15, 1, 1, 1
        if self.training:
            noise = torch.randn_like(route_1_raw) * stdev_1
        else:
            noise = 0

        route_1_raw = torch.softmax(route_1_raw + noise, dim=1)
        route_1_indice = torch.topk(route_1_raw, self.desert_k, dim=1, largest=False)[1]
        route_1 = route_1_raw.clone()
        for j in range(n):
            for i in range(self.desert_k):
                route_1[j, route_1_indice[j, i].reshape(-1)] = 0
        lora_out_1 = []
        for i in range(self.num_lora):
            lora_out_1.append(self.lora_list_1[i](out).unsqueeze(1)) # n, 1, c, h, w
        lora_out_1 = torch.cat(lora_out_1, dim=1)
        lora_out_1 = torch.sum(lora_out_1 * route_1, dim=1)
        out = self.bn_all[task](lora_out_1) + self.conv2[task](out) * prob_mix[:, 0] + self.share_conv(out.detach())
        out = self.bn[task](out)
        out = self.activate(out)

        out = self.conv3[task](out)
        return out, route_feat, route_1


class Residual(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn
    def forward(self, x):
        x1=self.fn(x)
        return x1+x

class DefMixer(nn.Module):
    def __init__(self,dim_in, dim, depth=1, kernel_size=1):
        super(DefMixer, self).__init__()

        self.blocks = nn.Sequential(
            *[nn.Sequential(
                    nn.Conv2d(dim, dim, kernel_size=kernel_size),
                    nn.GELU(),
                    nn.BatchNorm2d(dim),
                    Residual(nn.Sequential(
                        ChlSpl(dim, dim, (1, 3), 1, 0),
                        nn.GELU(),
                        nn.BatchNorm2d(dim)
                    )),
            ) for i in range(depth)],
        )

    def forward(self, x):
        x = self.blocks(x)
        return x

class ChlSpl(nn.Module):

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size,
        stride: int = 1,
        padding: int = 0,
        dilation: int = 1,
        groups: int = 1,
        bias: bool = True,
    ):
        super(ChlSpl, self).__init__()

        if in_channels % groups != 0:
            raise ValueError('in_channels must be divisible by groups')
        if out_channels % groups != 0:
            raise ValueError('out_channels must be divisible by groups')
        if stride != 1:
            raise ValueError('stride must be 1')
        if padding != 0:
            raise ValueError('padding must be 0')

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = _pair(stride)
        self.padding = _pair(padding)
        self.dilation = _pair(dilation)
        self.groups = groups

        self.weight = nn.Parameter(torch.empty(out_channels, in_channels // groups, 1, 1))

        self.get_offset = Offset(dim=in_channels, kernel_size=3)

        if bias:
            self.bias = nn.Parameter(torch.empty(out_channels))
        else:
            self.register_parameter('bias', None)
        self.register_buffer('offset', self.gen_offset())

        self.reset_parameters()

    def reset_parameters(self) -> None:
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))

        if self.bias is not None:
            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in)
            init.uniform_(self.bias, -bound, bound)
    def gen_offset(self):
        offset = torch.empty(1, self.in_channels*2, 1, 1)
        start_idx = (self.kernel_size[0] * self.kernel_size[1]) // 2
        assert self.kernel_size[0] == 1 or self.kernel_size[1] == 1, self.kernel_size
        for i in range(self.in_channels):
            if self.kernel_size[0] == 1:
                offset[0, 2 * i + 0, 0, 0] = 0
                offset[0, 2 * i + 1, 0, 0] = (i + start_idx) % self.kernel_size[1] - (self.kernel_size[1] // 2)
            else:
                offset[0, 2 * i + 0, 0, 0] = (i + start_idx) % self.kernel_size[0] - (self.kernel_size[0] // 2)
                offset[0, 2 * i + 1, 0, 0] = 0
        return offset

    def forward(self, input: Tensor) -> Tensor:
        """
            input: Tensor[b,c,h,w]
        """
        offset_2 = self.get_offset(input)
        B, C, H, W = input.size()

        return deform_conv2d_tv(input, offset_2, self.weight, self.bias, stride=self.stride, padding=self.padding, dilation=self.dilation)

    def extra_repr(self) -> str:
        s = self.__class__.__name__ + '('
        s += '{in_channels}'
        s += ', {out_channels}'
        s += ', kernel_size={kernel_size}'
        s += ', stride={stride}'
        s += ', padding={padding}' if self.padding != (0, 0) else ''
        s += ', dilation={dilation}' if self.dilation != (1, 1) else ''
        s += ', groups={groups}' if self.groups != 1 else ''
        s += ', bias=False' if self.bias is None else ''
        s += ')'
        return s.format(**self.__dict__)

class Offset(nn.Module):
    def __init__(self, dim, kernel_size):
        super().__init__()
        self.kernel_size = kernel_size
        self.stride = 1
        self.p_conv = nn.Conv2d(dim, 2*kernel_size*kernel_size, kernel_size=3, padding=1, stride=1)
        nn.init.constant_(self.p_conv.weight, 0)
        self.p_conv.register_backward_hook(self._set_lr)
        self.opt = nn.Conv2d(2*self.kernel_size*self.kernel_size, dim*2, kernel_size=3, padding=1, stride=1, groups=2)


    @staticmethod
    def _set_lr(module, grad_input, grad_output):
        grad_input = (grad_input[i] * 0.1 for i in range(len(grad_input)))
        grad_output = (grad_output[i] * 0.1 for i in range(len(grad_output)))

    def _get_p(self, offset, dtype):
        N, h, w = offset.size(1)//2, offset.size(2), offset.size(3)

        p_n = self._get_p_n(N, dtype)
        p_0 = self._get_p_0(h, w, N, dtype)
        p = p_0 + p_n + offset
        return p

    def _get_p_n(self, N, dtype):
        p_n_x, p_n_y = torch.meshgrid(
            torch.arange(-(self.kernel_size-1)//2, (self.kernel_size-1)//2+1),
            torch.arange(-(self.kernel_size-1)//2, (self.kernel_size-1)//2+1))
        p_n = torch.cat([torch.flatten(p_n_x), torch.flatten(p_n_y)], 0)
        p_n = p_n.view(1, 2*N, 1, 1).type(dtype)
        return p_n

    def _get_p_0(self, h, w, N, dtype):
        p_0_x, p_0_y = torch.meshgrid(
            torch.arange(1, h*self.stride+1, self.stride),
            torch.arange(1, w*self.stride+1, self.stride))
        p_0_x = torch.flatten(p_0_x).view(1, 1, h, w).repeat(1, N, 1, 1)
        p_0_y = torch.flatten(p_0_y).view(1, 1, h, w).repeat(1, N, 1, 1)
        p_0 = torch.cat([p_0_x, p_0_y], 1).type(dtype)
        return p_0

    def forward(self, x):
        offset = self.p_conv(x)
        dtype = offset.data.type()
        N = offset.size(1) // 2
        p = self._get_p(offset, dtype)  #1,18,107,140
        p =self.opt(p)
        return p

class ConvBNReLU(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, dilation=1, groups=1,
                 norm_layer=nn.BatchNorm2d, activation_layer=nn.ReLU, bias='auto',
                 inplace=True, affine=True):
        super().__init__()
        padding = dilation * (kernel_size - 1) // 2
        self.use_norm = norm_layer is not None
        self.use_activation = activation_layer is not None
        if bias == 'auto':
            bias = not self.use_norm
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding,
                              dilation=dilation, groups=groups, bias=bias)
        if self.use_norm:
            self.bn = norm_layer(out_channels, affine=affine)
        if self.use_activation:
            self.activation = activation_layer(inplace=inplace)

    def forward(self, x):
        x = self.conv(x)
        if self.use_norm:
            x = self.bn(x)
        if self.use_activation:
            x = self.activation(x)
        return x


